#pragma execution_character_set("utf-8")  

#include <windows.h>
#include <iostream>
#include <vector>
#include <string>
#include <fstream>
#include <gdiplus.h>
#include <shlobj.h>
#include <algorithm> // 包含 std::find_if 和 std::string::iterator
#include <cctype> // 包含 std::isspace
#include "rapidjson/document.h"
#include "rapidjson/prettywriter.h"

#pragma comment(lib, "gdiplus.lib")

using namespace Gdiplus;
using namespace std;

// 获取文件图标(大图标)
HICON GetFileIcon(LPCTSTR filePath) {
	SHFILEINFO sfi = { 0 };
	DWORD_PTR result = SHGetFileInfo(
		filePath,
		0,
		&sfi,
		sizeof(SHFILEINFO),
		SHGFI_USEFILEATTRIBUTES | SHGFI_ICON | SHGFI_LARGEICON
	);
	return (result != 0) ? sfi.hIcon : NULL;
}

// 获取文件夹图标(大图标)
HICON GetFolderIcon(LPCTSTR folderPath) {
	SHFILEINFO sfi = { 0 };
	DWORD_PTR result = SHGetFileInfo(
		folderPath,
		FILE_ATTRIBUTE_DIRECTORY,
		&sfi,
		sizeof(SHFILEINFO),
		SHGFI_USEFILEATTRIBUTES | SHGFI_ICON | SHGFI_LARGEICON
	);
	return (result != 0) ? sfi.hIcon : NULL;
}

// 获取文件夹图标(大图标)
HICON GetFolderIcon(int csidl) {
	SHFILEINFO sfi = { 0 };
	LPITEMIDLIST pidl;
	SHGetSpecialFolderLocation(NULL, csidl, &pidl);
	DWORD_PTR result = SHGetFileInfo(
		(LPCWSTR)pidl,
		FILE_ATTRIBUTE_DIRECTORY,
		&sfi,
		sizeof(SHFILEINFO),
		SHGFI_PIDL | SHGFI_ICON | SHGFI_LARGEICON
	);
	return (result != 0) ? sfi.hIcon : NULL;
}

std::string GetIconResourceFromDesktopIni(const std::string& folderPath) {
	std::string desktopIniPath = folderPath + "\\desktop.ini";
	std::ifstream file(desktopIniPath);
	std::string line;
	std::string iconResource;

	while (std::getline(file, line)) {
		if (line.find("[.ShellClassInfo]") != std::string::npos) {
			while (std::getline(file, line)) {
				if (line.find("IconResource") != std::string::npos) {
					size_t pos = line.find("=");
					if (pos != std::string::npos) {
						iconResource = line.substr(pos + 1);
						break;
					}
				}
				else if (line.empty()) {
					break;
				}
			}
			break;
		}
	}
	return iconResource;
}

HICON GetIconFromFile(LPCWSTR filePath, int index) {
	// 提取77号图标资源
	HICON hLargeIcon = NULL;
	HICON hSmallIcon = NULL;
	UINT extracted = ExtractIconExW(
		filePath,      // DLL路径
		index,           // 图标索引
		&hLargeIcon,  // 接收大图标
		&hSmallIcon,  // 接收小图标
		1             // 提取数量
	);
	return hLargeIcon;
}

LPCWSTR stringToLPCWSTR(const std::string& str) {
	int wideStrLen = MultiByteToWideChar(
		CP_ACP,                // 使用当前系统ANSI代码页
		0,                     // 无特殊标志
		str.c_str(),           // 输入字符串
		-1,                    // 自动计算长度
		NULL,                  // 不预分配缓冲区
		0                      // 请求计算所需缓冲区大小
	);

	if (wideStrLen == 0) return NULL;

	wchar_t* wideStr = new wchar_t[wideStrLen];
	MultiByteToWideChar(
		CP_ACP, 0,
		str.c_str(), -1,
		wideStr, wideStrLen
	);

	return wideStr;  // 注意：调用方需负责释放内存
}
#include <codecvt>
std::wstring StringToWstring(const std::string& str) {
	wstring_convert<std::codecvt_utf8<wchar_t>> converter;
	return converter.from_bytes(str);
}
int checkPathTypeWin(const std::string& path) {
	DWORD attrs = GetFileAttributesW(StringToWstring(path).c_str());
	if (attrs == INVALID_FILE_ATTRIBUTES) {
		return -1;
	}
	else if (attrs & FILE_ATTRIBUTE_DIRECTORY) {
		return 1;
	}
	else {
		return 2;
	}
}

std::string TCHARArrayToStdString(const TCHAR tcharArray[]) {
#ifdef _UNICODE
	int len = WideCharToMultiByte(CP_UTF8, 0, tcharArray, -1, NULL, 0, NULL, NULL);
	std::string result(len, 0);
	WideCharToMultiByte(CP_UTF8, 0, tcharArray, -1, &result[0], len, NULL, NULL);
	return result;
#else
	return std::string(tcharArray);
#endif
}

string GetPathByCSID(int csidl) {
	string sysPath;
	TCHAR szPath[MAX_PATH];
	LPITEMIDLIST pidl;
	SHGetSpecialFolderLocation(NULL, csidl, &pidl);
	if (SHGetPathFromIDList(pidl, szPath)) {
		sysPath = TCHARArrayToStdString(szPath);
	}
	return sysPath;
}

int GetCSIDByPath(const string path) {
	vector<int> CSIDS = { CSIDL_DESKTOP,CSIDL_INTERNET,CSIDL_PROGRAMS,CSIDL_CONTROLS,CSIDL_PRINTERS,CSIDL_PERSONAL,CSIDL_FAVORITES,CSIDL_STARTUP,CSIDL_RECENT,CSIDL_SENDTO,CSIDL_BITBUCKET,CSIDL_STARTMENU,CSIDL_MYDOCUMENTS,CSIDL_MYMUSIC,CSIDL_MYVIDEO,CSIDL_DESKTOPDIRECTORY,CSIDL_DRIVES,CSIDL_NETWORK,CSIDL_NETHOOD,CSIDL_FONTS,CSIDL_TEMPLATES,CSIDL_COMMON_STARTMENU,CSIDL_COMMON_PROGRAMS,CSIDL_COMMON_STARTUP,CSIDL_COMMON_DESKTOPDIRECTORY,CSIDL_APPDATA,CSIDL_PRINTHOOD };
	for (vector<int>::iterator it = CSIDS.begin(); it != CSIDS.end(); it++) {
		if (!strcmp(GetPathByCSID(*it).c_str(), path.c_str()))
			return *it;
	}
	return -1;
}

int GetEncoderClsid(const WCHAR* format, CLSID* pClsid) {
	UINT num = 0, size = 0;
	GetImageEncodersSize(&num, &size);
	if (size == 0) return -1;

	ImageCodecInfo* pCodecInfo = (ImageCodecInfo*)malloc(size);
	GetImageEncoders(num, size, pCodecInfo);

	for (UINT i = 0; i < num; ++i) {
		if (wcscmp(pCodecInfo[i].MimeType, format) == 0) {
			*pClsid = pCodecInfo[i].Clsid;
			free(pCodecInfo);
			return i;
		}
	}
	free(pCodecInfo);
	return -1;
}

std::string HICONToBase64(HICON hIcon, const WCHAR* format) {
	GdiplusStartupInput gdiplusStartupInput;
	ULONG_PTR gdiplusToken;
	GdiplusStartup(&gdiplusToken, &gdiplusStartupInput, NULL);

	Bitmap *bitmap = Bitmap::FromHICON(hIcon);
	IStream* pStream = NULL;
	CreateStreamOnHGlobal(NULL, TRUE, &pStream);

	CLSID clsid;
	GetEncoderClsid(format, &clsid);
	bitmap->Save(pStream, &clsid);

	HGLOBAL hGlobal = NULL;
	GetHGlobalFromStream(pStream, &hGlobal);
	LPVOID pData = GlobalLock(hGlobal);
	DWORD dataSize = GlobalSize(hGlobal);

	std::vector<BYTE> buffer(dataSize);
	memcpy(&buffer[0], pData, dataSize);
	   
	const char base64Table[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
	std::string result;
	for (size_t i = 0; i < buffer.size(); i += 3) {
		BYTE b1 = buffer[i];
		BYTE b2 = (i + 1 < buffer.size()) ? buffer[i + 1] : 0;
		BYTE b3 = (i + 2 < buffer.size()) ? buffer[i + 2] : 0;

		result += base64Table[b1 >> 2];
		result += base64Table[((b1 & 0x03) << 4) | (b2 >> 4)];
		result += (i + 1 < buffer.size()) ? base64Table[((b2 & 0x0F) << 2) | (b3 >> 6)] : '=';
		result += (i + 2 < buffer.size()) ? base64Table[b3 & 0x3F] : '=';
	}

	GlobalUnlock(hGlobal);
	pStream->Release();
	delete bitmap;
	GdiplusShutdown(gdiplusToken);
	return result;
}

string getIconBase64ByPath(string filePath) {
	string base64str;
	HICON hIcon = NULL;
	int type = checkPathTypeWin(filePath);
	if (type == 1) {
		int CSID = GetCSIDByPath(filePath);
		if (CSID > -1) {
			hIcon = GetFolderIcon(CSID);
		}
		else {
			string IconResource = GetIconResourceFromDesktopIni(filePath);
			if (IconResource == "") {
				hIcon = GetFolderIcon(stringToLPCWSTR(filePath));
			}
			else {
				string imgPath = IconResource.substr(0, IconResource.find(","));
				string index = IconResource.substr(IconResource.find(",") + 1);
				hIcon = GetIconFromFile(stringToLPCWSTR(imgPath), stoi(index));
			}
		}
	}
	if (type == 2) {
		hIcon = GetFileIcon(stringToLPCWSTR(filePath));
	}
	if (hIcon) {
		base64str = HICONToBase64(hIcon, L"image/png");
		DestroyIcon(hIcon);
	}
	return base64str;
}
bool verification(string path) {
	if (path == "") return false;
	return true;
}
int main(int argc, char* argv[]) {
	if (argc < 2) 
		return 0;
	string p = argv[1];
	if (p == "-o") {	// 同步输出
		string filePath = argv[2];
		string ret = getIconBase64ByPath(filePath);
		fprintf(stdout, ret.c_str());
	}
	if (p == "-x") { // 管道模式
		while (1) {
			string input;
			std::getline(std::cin, input);
			if (verification(input)) {
				rapidjson::Document document;
				if (!document.Parse(input.c_str()).HasParseError()) {
					string filePath;
					if (document["path"].IsString()) {
						filePath = (string)document["path"].GetString();
						std::string base64str = getIconBase64ByPath(filePath);
						document.AddMember("base64str", rapidjson::Value().SetString(base64str.c_str(), document.GetAllocator()), document.GetAllocator());
						rapidjson::StringBuffer buffer;
						rapidjson::Writer<rapidjson::StringBuffer> writer(buffer);
						document.Accept(writer); // 将doc的内容写入到buffer中
						fprintf(stdout, "%s\n", buffer.GetString());
					}
				}
			}
		}
	}
	
	return 0;
}
